﻿using System;
using System.Diagnostics.CodeAnalysis;
using System.Text;
using PhoneNetTools.Utilities;

namespace PhoneNetTools.Dns.Records
{
    /// <summary>
    /// A helper class that is able to extract strings from raw DNS messages.
    /// </summary>
    internal static class DnsStringBuilder
    {
        /// <summary>
        /// Extracts a string from raw DNS message data, starting at the given offset.
        /// The method is able to handle DNS string compression correctly.
        /// The offset is advanced to the next data after the string entry,
        /// correctly handling compressed or uncompressed string sizes.
        /// </summary>
        /// <param name="rawData">The raw data used to extract the string.</param>
        /// <param name="startIndex">The start index to start extracting at.</param>
        /// <returns>The extracted string.</returns>
        public static string ExtractString(byte[] rawData, ref int startIndex)
        {
            if (rawData == null)
            {
                throw new ArgumentNullException("rawData");
            }

            // initial data
            var sb = new StringBuilder();
            var prefix = rawData[startIndex];

            // to keep track of compression
            // => this is required to adjust the "global" offset
            // => as long as we do not have compressed text, we simply proceed with the offset,
            // but as soon as compression happened once, we don't move the offset anymore
            var compressionHappened = false;

            // go
            var currentIndex = startIndex;
            while (prefix > 0)
            {
                // extract part
                bool isPartCompressed;
                ExtractPart(prefix, rawData, ref currentIndex, out isPartCompressed, sb);

                // decide what to do with the global offset, depending on the compression flag
                if (!compressionHappened)
                {
                    if (isPartCompressed)
                    {
                        // first compression detected
                        // => proceed once more by the number of bytes the compression algorithm eats up (2)
                        startIndex += 2;
                        compressionHappened = true;
                    }
                    else
                    {
                        // no compression happened so far, the local and global offset are in sync
                        startIndex = currentIndex;
                    }
                }

                // get the next prefix to see what we need to do
                prefix = rawData[currentIndex];
            }

            // if no compression happened, we need to move the global index once more
            // to compensate for the 0-termination
            if (!compressionHappened)
            {
                startIndex++;
            }

            // and done
            return sb.ToString();
        }

        private static void ExtractPart(byte prefix, byte[] rawData, ref int currentIndex, out bool isCompressed, StringBuilder sb)
        {
            // add a point if this is not the first part
            if (sb.Length > 0)
            {
                sb.Append('.');
            }

            // check whether compression is used for this part,
            // adjust the current index accordingly if applicable,
            // and get the number of characters to extract
            isCompressed = false;
            var count = prefix;
            if ((prefix & 0xc0) == 0xc0)
            {
                // this is using the text "compression" feature of the protocol
                isCompressed = true;

                // get the actual string offset in raw data
                var temp = BigEndianBitConverter.ToUInt16(rawData, currentIndex);
                currentIndex = temp & ~(0xc0 << 8);

                // also adjust count
                count = rawData[currentIndex];
            }

            // now proceed to the actual start of the string
            currentIndex++;

            // extract the actual string
            for (var i = currentIndex; i < currentIndex + count; i++)
            {
                sb.Append((char)rawData[i]);
            }

            // increment the current index
            currentIndex += count;
        }
    }
}